from zhipuapi import *
from novelreader import *
from novelplayer import *
import streamlit as st


def render_ui():
    st.title("小说角色情景对话")
    txt_bytes = st.file_uploader("请上传小说文件或截选章节（支持txt/markdown格式）", type=[".txt", ".md"], accept_multiple_files=False)
    rounds = st.slider("对话轮数（请不要在对话生成中修改）", min_value=5, max_value=20, value=12)
    if txt_bytes:
        st.text("对话生成中...")
        r = read_novel(txt_bytes, limit=None)
        for output in generate_conversation(r, output_file_path='texts/conversations.txt', round_time=rounds):
            st.markdown(output.strip())
        st.text("The End.")


def test_program():
    with open('texts/oldman-sea.txt') as f:
        print("Generate conversations...")
        r = read_novel(f, limit=100)
        for output in generate_conversation(r, output_file_path='texts/conversations.txt', round_time=5):
            print(output)
        print("End.")


if __name__ == '__main__':
    render_ui()